--[[
    令牌桶限流

    self.bucket_dict_name 需要与 nginx 配置文件中 http 段配置的共享内存区域一致，如：

    http {
        lua_shared_dict my_limit_conn_store 100m;

        server {
            location / {
                ...
            }
        }
    }

    Author: mahaixing@gmail.com
]]

local assert = assert
--限流
local limit_req = require "resty.limit.req"
local utils = require "utils"

local _M = require("limiters.limiter"):new()

function _M:process_config()
    -- 速率
    self.bucket_rate = assert(self.conf["bucket_rate"], "Please set bucket_rate in bucket section!")
    -- 桶
    self.bucket_burst = assert(self.conf["bucket_burst"], "Please set bucket_burst in bucket section!")
    -- 桶针对 uri 还是客户端ip地址
    self.bucket_target = assert(self.conf["bucket_target"], "Please set bucket_target in bucket section!")
    self.bucket_target = string.lower(self.bucket_target)
    print(self.bucket_target)
    if (self.bucket_target ~= "ip" and self.bucket_target ~= "uri") then
        error("bucket_target error must be \"ip\" or \"uri\"")
    end
    
    if (utils.is_null(self.conf["bucket_dict_name"])) then
        self.bucket_dict_name = "limit_srvs"
    end

    self.bucket_rate = tonumber(self.bucket_rate)
    if self.bucket_rate == nil then
        utils.log("bucket_rate must a number!", ngx.ERR)
        error("bucket_rate must a number!")
    end

    self.bucket_burst = tonumber(self.bucket_burst)
    if self.bucket_burst == nil then
        utils.log("bucket_burst must a number!", ngx.ERR)
        error("bucket_burst must a number!")
    end

end

function _M:execute()
    -- print (self.bucket_dict_name .. " " .. self.bucket_rate .. " " .. self.bucket_burst)

    local lim, err = limit_req.new(self.bucket_dict_name, self.bucket_rate, self.bucket_burst)
    if not lim then
        utils.log("Initialization limit_req error: " .. err, ngx.ERR)
        error("Initialization limit_req error: " .. err, ngx.ERR)
        -- 初始化失败，不限流
        -- return false
        return ngx.exit(500)
    end
    
    local key
    if self.bucket_target == "uri" then
        key = ngx.var.uri 
    else 
        key = ngx.var.binary_remote_addr
    end
    
    print ("key : " .. key)
    local delay, err = lim:incoming(key, true)

    print ("delay: " .. delay)
    -- 超出桶大小
    if not delay then
        if err == "rejected" then
            utils.log("rejected", ngx.INFO)
            return true
        end
        utils.log("failed to limit req: " .. err, ngx.INFO)
    end

    if delay > 0 then
        ngx.sleep(delay)
        utils.log("nginx sleep " .. delay .. " seconds", ngx.INFO)
    end

    return false
end

return _M